from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer

# 1. 初始化BPE分词器
tokenizer = Tokenizer(BPE(unk_token="<unk>"))

# 2. 准备训练器
trainer = BpeTrainer()

# 3. 训练分词器(需要提供文本文件)
files = ["./data/timemachine.txt"]  # 替换为你的文本文件
tokenizer.train(files, trainer)

# 4. 保存和加载模型
tokenizer.save("bpe_tokenizer.json") # 保存
tokenizer = Tokenizer.from_file("bpe_tokenizer.json") # 加载

# 5. 使用分词器
output = tokenizer.encode("Hello world!") # 对指定内容进行BPE分词
print(output.tokens)  # 输出分词结果
print(output.ids)  # 输出token id
